package com.tianji.aigc.controller;

import cn.hutool.core.collection.CollStreamUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;

import java.util.List;

@RequiredArgsConstructor
@RestController
@Slf4j
@RequestMapping("/embedding")
public class EmbeddingController {

    private final VectorStore vectorStore;
    private final EmbeddingModel embeddingModel;

    /**
     * 添加向量
     * @param messages
     */
    @PostMapping
    public void addEmbedding(@RequestParam List<String> messages) {
        log.info("保存到向量数据库中，消息数据：{}", messages);
        List<Document> documents = CollStreamUtil.toList(messages, message -> Document.builder().text(message).build());
        vectorStore.add(documents);
        log.info("保存到向量数据库成功, 数量：{}", messages.size());
    }

    /**
     * 文本转换为向量
     * @param message
     * @return
     */
    @GetMapping
    public EmbeddingResponse tranceEmbedding(@RequestParam("message") String message) {
        log.info("开始向量转换");
        return this.embeddingModel.embedForResponse(List.of(message));
    }

    /**
     * 向量查询
     */
    @GetMapping("/search")
    public List<Document> search(@RequestParam("message") String message) {
        log.info("开始向量查询");
        return vectorStore.similaritySearch(
                SearchRequest
                        .builder()
                        .query(message)
                        .similarityThreshold(0.6d)
                        .topK(3)
                        .build()
                );
    }

    /**
     * 全部向量查询
     */
    @GetMapping("/search/all")
    public List<Document> searchAll() {
        log.info("开始全部向量查询");
        return vectorStore.similaritySearch(SearchRequest.builder().query("").topK(999).build());
    }

    /**
     * 删除向量
     */
    @DeleteMapping
    public void deleteByIds(@RequestParam List<String> ids) {
        log.info("开始删除向量");
        vectorStore.delete(ids);
    }


}
